{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  OchreLang.ExpandMain
-- Copyright   :  (c) Stephen Tetley 2014
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Inline all method calls in main.
--
-- Works on intermediate syntax because we need alpha renaming.
--
--------------------------------------------------------------------------------

module OchreLang.ExpandMain
  ( 
    expandMain

  ) where

import qualified OchreLang.Alpha as Alpha
import OchreLang.CompilerMon
import OchreLang.Subst
import OchreLang.Syntax
import OchreLang.SyntaxBase


import Control.Applicative
import qualified Data.Map as Map


-- | Super-inlining has already taken place...
--
type LocalMethodDict = Map.Map Ident MethodDefn
type Env             = LocalMethodDict


findDefn :: LocalMethodDict -> Ident -> Maybe MethodDefn
findDefn = flip Map.lookup


buildLocalMethodDict :: [MethodDecl] -> LocalMethodDict
buildLocalMethodDict = foldr fn Map.empty
  where
    fn (MethodDecl pid defn) = Map.insert pid defn


-- Don\'t initNameGen
expandMain :: Instrument -> Compiler Instrument
expandMain inst@(Instrument { instr_main       = main_expr
                            , instr_methods    = mdecls }) =  
    do { main1 <- inline (buildLocalMethodDict mdecls) main_expr
       ; return $ inst { instr_main       = main1
                       , instr_methods    = [] 
                       }
       }




-- | There is no need to track productivity as inlining proceeds
-- by recursive descent.
--
inline :: Env -> Expr -> Compiler Expr
inline = inlineExpr



inlineExpr :: Env -> Expr -> Compiler Expr
inlineExpr _   (Return val)                 = return $ Return val

inlineExpr env (OpCall vid vs)              = case findDefn env vid of
    Nothing -> return $ OpCall vid vs
    Just (MethodDefn args body) -> do { senv <- insertsArgs args vs Map.empty
                                      ; e1   <- Alpha.alphaK_Expr =<< substArgs senv body
                                      ; inlineExpr env e1
                                      }

inlineExpr _   (SuperCall {})               = 
    throwError "ExpandMain.inlineExpr" "Super-Call not inlined before expand-main"

inlineExpr _   (FunApp {})                  = 
    throwError "ExpandMain.inlineExpr" "Fun-App not inlined before expand-main"
  
inlineExpr env (LetValue vid val e)         = 
    LetValue vid val <$> inlineExpr env e

inlineExpr env (LetSig1 vid e1 e2)          = 
    LetSig1 vid <$> inlineExpr env e1 <*> inlineExpr env e2

inlineExpr env (LetSig2 v1 v2 e1 e2)        = 
    LetSig2 v1 v2 <$> inlineExpr env e1 <*> inlineExpr env e2

inlineExpr env (LetSig3 v1 v2 v3 e1 e2)     = 
    LetSig3 v1 v2 v3 <$> inlineExpr env e1 <*> inlineExpr env e2

inlineExpr env (LetSig4 v1 v2 v3 v4 e1 e2)  = 
    LetSig4 v1 v2 v3 v4 <$> inlineExpr env e1 <*> inlineExpr env e2

inlineExpr env (LetTuple vs e1 e2)          = 
    LetTuple vs <$> inlineExpr env e1 <*> inlineExpr env e2

inlineExpr _   (LetFun {})                  = 
    throwError "ExpandMain.inlineExpr" "Let-Fun not inlined before expand-main"

inlineExpr env (Ifte val e1 e2)             =
    Ifte val <$> inlineExpr env e1 <*> inlineExpr env e2

inlineExpr env (e1 :>> e2)                  =
    (:>>) <$> inlineExpr env e1 <*> inlineExpr env e2



